feat(token-id-capture): deliver rebuilt trajectories safely by ananthsub · Pull Request #2126 · NVIDIA-NeMo/Gym

ananthsub · 2026-07-23T15:35:23Z

Attaches a safely rebuilt token trajectory to each participating rollout and retires capture evidence only after durable handoff. This PR also owns the durable capture lifecycle used by the later lineage and prefix-supply PRs.

Capture custody and delivery

sequenceDiagram
    participant M as Model worker
    participant S as TokenSink / TokenSource backend
    participant C as Rollout collector or framework
    participant B as Trajectory builder
    participant D as Durable downstream

    M->>S: begin_call(rollout_id, model_call_id) when supported
    Note over M,S: Failure occurs before generation - a dangling intent identifies a lost entry
    M->>M: run inference and build TokenEntry
    M->>S: await put(TokenEntry)
    S-->>M: entry durable
    M-->>M: return model response

    C->>S: freeze(rollout_id)
    S-->>C: TokenCaptureSnapshot(entries, incomplete, snapshot_id, version)
    C->>B: rebuild frozen snapshot
    alt incomplete, malformed, ambiguous, split, or empty
        B-->>C: mask_sample=true and diagnostics
        Note over C,S: retain evidence
    else one safe trainable chain
        B-->>C: rebuilt response and metrics
        C->>D: write rebuilt rollout
        D-->>C: durable acknowledgement
        C->>S: drop(rollout_id, snapshot_id, version)
        S-->>C: true only if the frozen version is still current
    end

A write racing freeze may become durable, but it advances the version so a stale drop cannot erase it. Retirement keeps a tombstone until operator-controlled garbage collection, preventing an old worker from recreating a retired rollout.

Summary

Clears and finalizes records only for agents selected by static token_id_capture configuration or all_agents; tokenless nonparticipants remain ordinary evaluation rollouts.
Defines TokenSink.put as an awaited durability boundary and supports optional pre-dispatch begin_call custody. A dangling intent makes the frozen snapshot incomplete when both put and mark_incomplete fail.
Freezes entries and incomplete state as one versioned snapshot, deduplicates at-least-once entries, and conditionally retires only the consumed snapshot_id and version after downstream durability.
Masks missing, malformed, incomplete, ambiguous, multi-root, multi-chain, and empty-token builds. max_mask_fraction can stop a run that produces too many masked rollouts.
Recovers a lagging state index from the durable JSONL tail, repairs an unacknowledged torn final line, reduces append-path state fsyncs, and provides sweep_retired for explicit tombstone cleanup.
Fsyncs Gym's JSONL result before retirement and retains failed or masked evidence.
Uses an installed caller-owned source without closing it, or constructs and closes Gym's default file source.
Fails before rollout dispatch when selected rows require rebuilding and the collector process has no source.
Provides transport conformance checks for publication, incomplete state, freeze, idempotency, conditional retirement, and optional call-intent custody.

Depends on #2125. The integration guide begins in #2341.

copy-pr-bot · 2026-07-23T15:35:26Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

copy-pr-bot · 2026-07-29T06:55:15Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

cmunley1

left a question above, but i think its unlikely in normal cases and we could follow up unless you want to reply there now.

ananthsub · 2026-08-18T05:30:36Z

/claude review

claude · 2026-08-18T05:40:54Z

SHIP — no reliability concerns.

Reviewed the token-capture delivery step and its wiring into run_rollouts. The correctness-critical properties hold:

Durable-before-delete ordering is correct. finalize_rollout_token_capture mutates result["response"]["output"] before orjson.dumps(result), so the persisted record carries the rebuilt trajectory; os.fsync(results_file.fileno()) runs before retire_rollout_token_capture, so source records are dropped only after the rebuilt trajectory is durable. No crash-window data loss.
Reward/scorer untouched — only response.output is replaced; reward and components are left alone (asserted in tests).
Never-raises holds — trajectories_from_source and source.drop are both wrapped; a failing transport masks one rollout, not the batch.
Retirement is safe against deterministic-id reuse — drop is conditioned on seal_id+version; a stale build cannot delete a newer attempt's records. Masked/failed builds are retained as evidence, not retired.
Ownership/lifecycle correct — owned_token_source is closed only when Gym constructed it (configured source or file store), never a process-installed source; second finalize call is a no-op.
Async hygiene clean — no httpx, no ray.get, all coroutines awaited; no undeclared imports; all symbols exported from the package __init__.

Test coverage is real and exercises masking, no-capture-key, nothing-recorded, mixed native/external batches, retention, and stale-record clearing.

NOTE (non-blocking, author's call): finalize_rollout_token_capture is invoked for every result including no_persist (kill_shaped) and failure_class rows that are never persisted, so a rebuild + possible masking warning is computed and discarded for those. Harmless — no reward/output impact and results aren't written — just minor wasted work and log noise on failure-heavy runs.

ananthsub · 2026-08-18T07:52:24Z

/claude review

claude · 2026-08-18T08:01:44Z

Review incomplete: the PR revision changed during review. The captured base (970603d) and head (84755ea) that this review analyzed no longer match the live PR, which now points at base b9746c7 and head 21e7ab0. The final pre-publish revision check therefore failed, so inline findings are withheld to avoid anchoring comments to a stale commit. Please re-run the review against the current head.

ananthsub · 2026-08-18T08:12:55Z

/claude review

cmunley1 · 2026-08-20T05:34:47Z

+    An unusable build has no rebuilt response and sets ``mask_sample``.
+    """
+    if source is None or rollout_carries_token_ids(result):
+        return None


With all_agents=True, if capture records are written, this early return skips the snapshot retirement https://github.com/NVIDIA-NeMo/Gym/pull/2126/changes#diff-4ed3026c5b1fe1967854548a34299c33ba36c8fd55c7cd2249324108dd0a550cR915, so those redundant records may accumulate. Should we avoid capture for agents who provide the token ids or still retire their captured snapshot?

A rollout's records are read, built into one contiguous response, and swapped in for response.output, so a trainer reads response.output the same way for a native agent and an external harness. The step takes the record and a TokenSource, and nothing else. Both belong to the caller, so every configuration question stays on the side that owns it: Gym's rollout collection resolves its source from Gym's config, and a training framework staging records through its own data plane passes its own source and never reads Gym's config. That is what makes the step shared rather than Gym's alone; a caller driving run_examples gets the same behaviour without reimplementing it. What a rollout needs is read from the rollout. An item already carrying token ids holds what the policy sampled, so that rollout is left alone; without them, the recorded ids are attached. A batch can therefore mix native agents and external harnesses, and a native agent moving onto capture simply stops arriving with ids and starts being rebuilt, with no config change anywhere. It also makes the step idempotent, since a rebuilt rollout carries ids. Records are cleared before dispatch, because rollout ids are deterministic and a store appends, and retired through the source once consumed. A rollout that lost a call is masked rather than trained on with a hole, reported as one field, mask_sample, at the top of the record: a consumer decides by reading it and nothing else. The metrics dict keeps the reasons, not the verdict. token_id_capture.rebuild_response says whether Gym reads records back at all. Off, Gym stops after the write and the caller reads through its own source; without it, a run capturing to its own transport has Gym reading a store nothing wrote to. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Resolve paired sources for readback and delay conditional capture retirement until the rebuilt rollout is fsynced, preserving masked evidence and newer attempts. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Use lifecycle terminology that states when rollout records become immutable without implying that durable delivery or retirement already completed. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

State the freeze, rebuild, handoff, and conditional retirement boundaries directly. Keep retention and masking comments aligned with the implementation. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Use the same static agent selection as request routing when clearing and rebuilding capture records. Tokenless agents that did not opt in remain ordinary evaluation rollouts. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Use caller-owned installed sources without closing them and fail before dispatch when selected rollouts cannot be rebuilt. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Record pre-dispatch call intent and preserve evidence when capture or retirement cannot complete, so consumers mask incomplete rollouts instead of training through silent holes. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

ananthsub changed the title ~~ananthsub/tokidcap/delivery~~ feat(token-id-capture): uniform delivery, per-agent scoping, retention Jul 23, 2026

ananthsub marked this pull request as ready for review July 23, 2026 16:18

copy-pr-bot Bot temporarily deployed to public July 23, 2026 16:19 Inactive

copy-pr-bot Bot temporarily deployed to public July 23, 2026 16:20 Inactive

copy-pr-bot Bot temporarily deployed to public July 23, 2026 16:21 Inactive

ananthsub requested review from bxyu-nvidia and cmunley1 July 23, 2026 16:22

github-actions Bot added the sla:review-overdue Review response is over the one-business-day SLA label Jul 24, 2026

github-actions Bot mentioned this pull request Jul 24, 2026

[Tracker] Pull request handoff SLA #2067

Open

ananthsub force-pushed the ananthsub/tokidcap/delivery branch from 74a48d6 to 71de7c5 Compare July 29, 2026 06:55

ananthsub force-pushed the ananthsub/tokidcap/delivery branch from 71de7c5 to 5c317c1 Compare July 29, 2026 07:12

ananthsub marked this pull request as draft July 29, 2026 07:12

github-actions Bot removed the sla:review-overdue Review response is over the one-business-day SLA label Jul 29, 2026

ananthsub force-pushed the ananthsub/tokidcap/delivery branch from 5c317c1 to 58cde9d Compare July 29, 2026 12:38

ananthsub force-pushed the ananthsub/tokidcap/delivery branch from f5fa8c6 to 06bb256 Compare August 5, 2026 15:57

copy-pr-bot Bot temporarily deployed to public August 5, 2026 15:57 Inactive

copy-pr-bot Bot temporarily deployed to public August 5, 2026 15:59 Inactive

copy-pr-bot Bot temporarily deployed to public August 5, 2026 16:00 Inactive

ananthsub force-pushed the ananthsub/tokidcap/delivery branch from 06bb256 to a28b020 Compare August 5, 2026 16:24

cmunley1 reviewed Aug 6, 2026

View reviewed changes

Comment thread nemo_gym/token_id_capture/delivery.py Outdated

cmunley1 previously approved these changes Aug 6, 2026

View reviewed changes

cmunley1 reviewed Aug 20, 2026

View reviewed changes

cmunley1 previously approved these changes Aug 20, 2026

View reviewed changes

cmunley1 previously approved these changes Aug 21, 2026

View reviewed changes

kajalj22 previously approved these changes Aug 21, 2026

View reviewed changes

cmunley1 previously approved these changes Aug 21, 2026

View reviewed changes

kajalj22 approved these changes Aug 21, 2026

View reviewed changes

ananthsub added 8 commits August 21, 2026 09:33

fix(token-id-capture): retire only durable snapshots

10ba544

Resolve paired sources for readback and delay conditional capture retirement until the rebuilt rollout is fsynced, preserving masked evidence and newer attempts. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

refactor(token-id-capture): freeze delivery snapshots

65a3e21

Use lifecycle terminology that states when rollout records become immutable without implying that durable delivery or retirement already completed. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

docs(token-id-capture): clarify durable delivery lifecycle

4b70d54

State the freeze, rebuild, handoff, and conditional retirement boundaries directly. Keep retention and masking comments aligned with the implementation. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

fix(token-id-capture): resolve sources in the collector

351c302

Use caller-owned installed sources without closing them and fail before dispatch when selected rollouts cannot be rebuilt. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

fix(token-id-capture): retire redundant native captures

fb633ba

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

fix(token-id-capture): harden durable delivery

247167e

Record pre-dispatch call intent and preserve evidence when capture or retirement cannot complete, so consumers mask incomplete rollouts instead of training through silent holes. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

cwing-nvidia mentioned this pull request Aug 22, 2026

chore(hermes_agent): evaluate replacing the fork commit with a named upstream release #2672

Open

7 tasks

Navigation Menu

Uh oh!

feat(token-id-capture): deliver rebuilt trajectories safely - #2126

feat(token-id-capture): deliver rebuilt trajectories safely#2126
ananthsub merged 8 commits into
mainfrom
ananthsub/tokidcap/delivery

ananthsub commented Jul 23, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jul 23, 2026

Uh oh!

copy-pr-bot Bot commented Jul 29, 2026

Uh oh!

Uh oh!

cmunley1 left a comment

Uh oh!

ananthsub commented Aug 18, 2026

Uh oh!

claude Bot commented Aug 18, 2026

Uh oh!

ananthsub commented Aug 18, 2026

Uh oh!

claude Bot commented Aug 18, 2026

Uh oh!

ananthsub commented Aug 18, 2026

Uh oh!

cmunley1 Aug 20, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

ananthsub commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Capture custody and delivery

Summary

Uh oh!

copy-pr-bot Bot commented Jul 23, 2026

Uh oh!

copy-pr-bot Bot commented Jul 29, 2026

Uh oh!

Uh oh!

cmunley1 left a comment

Choose a reason for hiding this comment

Uh oh!

ananthsub commented Aug 18, 2026

Uh oh!

claude Bot commented Aug 18, 2026

Uh oh!

ananthsub commented Aug 18, 2026

Uh oh!

claude Bot commented Aug 18, 2026

Uh oh!

ananthsub commented Aug 18, 2026

Uh oh!

cmunley1 Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ananthsub commented Jul 23, 2026 •

edited

Loading

cmunley1 Aug 20, 2026 •

edited

Loading